const std = @import("std");

const bpf = @import("bpf");
const Xdp = bpf.Xdp;
const BPF = std.os.linux.BPF;
const helpers = BPF.kern.helpers;
const trace_printk = helpers.trace_printk;

const __u8 = u8;
const __be16 = u16;
const __be32 = u32;
const __sum16 = u16;
const __u16 = u16;

const EthHdr = extern struct {
    dest: [6]u8,
    src: [6]u8,
    proto: u16,
};

const IPv4Hdr = extern struct {
    ver_ihl: u8,
    tos: u8,
    tot_len: u16,
    id: u16,
    frag_off: u16,
    ttl: u8,
    proto: u8,
    check: u16,
    src: u32,
    dst: u32,
};

const IPv6Hdr = extern struct {
    flow: u32,
    plen: u16,
    nxt: u8,
    hlim: u8,
    src: [16]u8,
    dst: [16]u8,
};

const TcpHdr = packed struct {
    source: __be16,
    dest: __be16,
    seq: __be32,
    ack_seq: __be32,
    res1: u4,
    doff: u4,
    fin: u1,
    syn: u1,
    rst: u1,
    psh: u1,
    ack: u1,
    urg: u1,
    ece: u1,
    cwr: u1,
    window: __be16,
    check: __sum16,
    urg_ptr: __be16,
};

const UdpHdr = extern struct {
    source: __be16,
    dest: __be16,
    len: __be16,
    check: __sum16,
};

const In6Addr = extern struct {
    in6_u: extern union {
        u6_addr8: [16]__u8,
        u6_addr16: [8]__be16,
        u6_addr32: [4]__be32,
    },
};

const ETH_PROTO_IPV4 = 0x0800;
const ETH_PROTO_IPV6 = 0x86DD;

const IPPROTO_ICMP = 1;
const IPPROTO_ICMPV6 = 58;
const IPPROTO_TCP = 6;
const IPPROTO_UDP = 17;

const BITMAP_TYPE = u64;

const Ipv4LpmKey = extern struct {
    prefixlen: u32,
    data: u32,
};

var src_map = bpf.Map.LpmTrie("src_map", Ipv4LpmKey, BITMAP_TYPE, 64, 1).init();

var dport_map = bpf.Map.HashMap("dport_map", u16, BITMAP_TYPE, 65535, 1).init();

var proto_map = bpf.Map.HashMap("proto_map", u8, BITMAP_TYPE, 256, 1).init();

var action_map = bpf.Map.HashMap("action_map", u64, c_int, 64, 1).init();

const DEBUG = true;

export fn firewall(ctx: *Xdp.Meta) linksection("xdp") c_int {
    const eth_hdr: *const EthHdr = ctx.get_ptr(EthHdr, 0) orelse return @intFromEnum(Xdp.RET.drop);

    const ret = switch (eth_hdr.proto) {
        std.mem.nativeTo(u16, ETH_PROTO_IPV4, .big) => handle_ipv4(ctx),
        std.mem.nativeTo(u16, ETH_PROTO_IPV6, .big) => handle_ipv6(ctx),
        else => {
            return @intFromEnum(Xdp.RET.aborted);
        },
    };
    return ret;
}

fn check(saddr: u32, dport: u16, proto: u8) c_int {
    var bitmap_array: [3]u64 = [_]u64{ 0, 0, 0 };
    var bitmap_array_len: u32 = 0;

    const src_key = Ipv4LpmKey{
        .prefixlen = 32,
        .data = saddr,
    };

    if (DEBUG) {
        const fmt = "saddr: %u";
        _ = trace_printk(fmt, fmt.len + 1, saddr, 0, 0);
    }

    const src_bitmap: ?*u64 = src_map.lookup(src_key);
    if (src_bitmap) |val| {
        if (DEBUG) {
            const fmt = "src_bitmap: %llu";
            _ = trace_printk(fmt, fmt.len + 1, val.*, 0, 0);
        }

        bitmap_array[bitmap_array_len] = val.*;
        bitmap_array_len += 1;
    }

    if (DEBUG) {
        const fmt = "dport: %u";
        _ = trace_printk(fmt, fmt.len + 1, dport, 0, 0);
    }

    const dport_key = dport;
    const dport_bitmap: ?*u64 = dport_map.lookup(dport_key);
    if (dport_bitmap) |val| {
        if (DEBUG) {
            const fmt = "dport_bitmap: %llu";
            _ = trace_printk(fmt, fmt.len + 1, val.*, 0, 0);
        }

        bitmap_array[bitmap_array_len] = val.*;
        bitmap_array_len += 1;
    }

    if (DEBUG) {
        const fmt = "proto: %u";
        _ = trace_printk(fmt, fmt.len + 1, proto, 0, 0);
    }

    const proto_key = proto;
    const proto_bitmap: ?*u64 = proto_map.lookup(proto_key);
    if (proto_bitmap) |val| {
        if (DEBUG) {
            const fmt = "proto_bitmap: %llu";
            _ = trace_printk(fmt, fmt.len + 1, val.*, 0, 0);
        }

        bitmap_array[bitmap_array_len] = val.*;
        bitmap_array_len += 1;
    }

    var result: u64 = ~@as(u64, 0);
    for (bitmap_array, 0..bitmap_array_len) |value, i| {
        if (DEBUG) {
            const fmt = "bitmap_array[%d] = %llu";
            _ = trace_printk(fmt, fmt.len + 1, i, value, 0);
        }

        result &= value;
    }

    if (DEBUG) {
        const fmt = "result1 = %llu";
        _ = trace_printk(fmt, fmt.len + 1, result, 0, 0);
    }

    var result2: i64 = @intCast(result);
    result2 &= -result2;

    if (DEBUG) {
        const fmt = "result2 = %llu";
        _ = trace_printk(fmt, fmt.len + 1, result, 0, 0);
    }

    const action_key = result;
    const action: ?*c_int = action_map.lookup(action_key);
    if (action) |val| {
        if (DEBUG) {
            const fmt = "action = %d";
            _ = trace_printk(fmt, fmt.len + 1, @intCast(val.*), 0, 0);
        }

        return val.*;
    } else {
        return @intFromEnum(Xdp.RET.drop);
    }
}

fn handle_ipv4(ctx: *Xdp.Meta) c_int {
    const iphdr_offset = @sizeOf(EthHdr);
    const ipv4hdr: *const IPv4Hdr = ctx.get_ptr(IPv4Hdr, iphdr_offset) orelse return @intFromEnum(Xdp.RET.aborted);

    if (ipv4hdr.proto == IPPROTO_ICMP) {
        return @intFromEnum(Xdp.RET.drop);
    } else if (ipv4hdr.proto == IPPROTO_TCP) {
        const tcpHdr: *const TcpHdr = ctx.get_ptr(TcpHdr, iphdr_offset + (ipv4hdr.ver_ihl * 4)) orelse return @intFromEnum(Xdp.RET.aborted);

        if (!(tcpHdr.syn == 1 and tcpHdr.ack == 0)) {
            return @intFromEnum(Xdp.RET.pass);
        }

        return check(ipv4hdr.src, tcpHdr.dest, ipv4hdr.proto);
    } else if (ipv4hdr.proto == IPPROTO_UDP) {
        const udpHdr: *const UdpHdr = ctx.get_ptr(UdpHdr, iphdr_offset + (ipv4hdr.ver_ihl * 4)) orelse return @intFromEnum(Xdp.RET.aborted);
        return check(ipv4hdr.src, udpHdr.dest, ipv4hdr.proto);
    } else {
        return @intFromEnum(Xdp.RET.drop);
    }
}

fn handle_ipv6(ctx: *Xdp.Meta) c_int {
    const iphdr_offset = @sizeOf(EthHdr);
    const ipv6hdr: *const IPv6Hdr = ctx.get_ptr(IPv6Hdr, iphdr_offset) orelse return @intFromEnum(Xdp.RET.aborted);

    if (ipv6hdr.nxt == IPPROTO_ICMPV6) {
        return @intFromEnum(Xdp.RET.pass);
    } else if (ipv6hdr.nxt == IPPROTO_TCP) {
        const tcpHdr: *const TcpHdr = ctx.get_ptr(TcpHdr, iphdr_offset + (10 * 4)) orelse return @intFromEnum(Xdp.RET.aborted);
        if (!(tcpHdr.syn == 1 and tcpHdr.ack == 0)) {
            return @intFromEnum(Xdp.RET.pass);
        }

        const fmt = "dport %u\n";
        _ = trace_printk(fmt, fmt.len + 1, tcpHdr.dest, 0, 0);
        return @intFromEnum(Xdp.RET.drop);
    } else if (ipv6hdr.nxt == IPPROTO_UDP) {
        const udpHdr: *const UdpHdr = ctx.get_ptr(UdpHdr, iphdr_offset + (10 * 4)) orelse return @intFromEnum(Xdp.RET.aborted);

        const fmt = "dport %u\n";
        _ = trace_printk(fmt, fmt.len + 1, udpHdr.dest, 0, 0);
        return @intFromEnum(Xdp.RET.drop);
    } else {
        return @intFromEnum(Xdp.RET.drop);
    }
}
